import torch
import cv2
import numpy as np
import os
import time

from App.model.plate_recognition.plateNet import myNet_ocr_color


def cv_imread(path):  #可以读取中文路径的图片
    img=cv2.imdecode(np.fromfile(path,dtype=np.uint8),-1)
    return img

def allFilePath(rootPath,allFIleList):
    fileList = os.listdir(rootPath)
    for temp in fileList:
        if os.path.isfile(os.path.join(rootPath,temp)):
            if temp.endswith('.jpg') or temp.endswith('.png') or temp.endswith('.JPG'):
                allFIleList.append(os.path.join(rootPath,temp))
        else:
            allFilePath(os.path.join(rootPath,temp),allFIleList)
device = torch.device('cuda') if torch.cuda.is_available() else torch.device("cpu")

color=['黑色','蓝色','绿色','白色','黄色']

plateName = r"#京沪津渝冀晋蒙辽吉黑苏浙皖闽赣鲁豫鄂湘粤桂琼川贵云藏陕甘青宁新学警港澳挂使领民航危0123456789ABCDEFGHJKLMNPQRSTUVWXYZ险品"
mean_value,std_value=(0.588,0.193)
def decodePlate(preds):
    pre=0
    newPreds=[]
    index=[]
    for i in range(len(preds)):
        if preds[i]!=0 and preds[i]!=pre:
            newPreds.append(preds[i])
            index.append(i)
        pre=preds[i]
    return newPreds,index

def image_processing(img,device):
    img = cv2.resize(img, (168,48))
    img = np.reshape(img, (48, 168, 3))

    # normalize
    img = img.astype(np.float32)
    img = (img / 255. - mean_value) / std_value
    img = img.transpose([2, 0, 1])
    img = torch.from_numpy(img)

    img = img.to(device)
    img = img.view(1, *img.size())
    return img
def get_plate_result(img,device,model,is_color=False):
    input = image_processing(img,device)
    if is_color:  #是否识别颜色
        preds,color_preds = model(input)
        color_preds = torch.softmax(color_preds,dim=-1)
        color_conf,color_index = torch.max(color_preds,dim=-1)
        color_conf=color_conf.item()
    else:
        preds = model(input)
    preds=torch.softmax(preds,dim=-1)
    prob,index=preds.max(dim=-1)
    index = index.view(-1).detach().cpu().numpy()
    prob=prob.view(-1).detach().cpu().numpy()
   
    
    # preds=preds.view(-1).detach().cpu().numpy()
    newPreds,new_index=decodePlate(index)
    prob=prob[new_index]
    plate=""
    for i in newPreds:
        plate+=plateName[i]
    # if not (plate[0] in plateName[1:44] ):
    #     return ""
    if is_color:
        return plate,prob,color[color_index],color_conf    #返回车牌号以及每个字符的概率,以及颜色，和颜色的概率
    else:
        return plate,prob
'''
device: 指定模型加载到哪个设备上运行（如 CPU 或 GPU）。通常通过 torch.device("cuda")
model_path: 是预训练模型权重文件的路径，该权重文件将用于加载训练好的模型。
is_color: 一个布尔参数，表示是否加载车牌颜色分类模型。默认为 False，表示只进行字符识别；如果为 True，则同时进行颜色分类。
'''
def init_model(device,model_path,is_color = False):


    # 加载车牌模型
    check_point = torch.load(model_path,map_location = device)
    # state_dict: 从检查点中提取模型的权重参数。这是一个包含模型各层参数的字典，用于恢复模型的权重状态。
    model_state= check_point['state_dict']
    # cfg: 从检查点中提取配置文件，包含了模型结构相关的配置信息。模型初始化时，可能需要根据这个配置文件来构建相同的网络架构。
    cfg = check_point['cfg']
    color_classes=0
    # 如果需要识别车牌颜色，进行分类
    if is_color:
        #颜色类别数
        color_classes=5
    #根据配置文件和指定的颜色类别数，创建一个 OCR 和颜色识别结合的模型。
    model = myNet_ocr_color(num_classes=len(plateName),export=True,cfg=cfg,color_num=color_classes)
    # 将提取的模型权重加载到初始化的模型中
    model.load_state_dict(model_state,strict=False)
    # 将模型分配到指定的设备（如 GPU 或 CPU）上，以便利用设备的计算能力进行推理。
    model.to(device)
    # 将模型设置为评估模式（Evaluation Mode），此时模型中的层如 dropout 和 batch normalization 将以推理模式运行，而不是训练模式。
    model.eval()
    # 返回经过初始化、加载权重、分配设备并处于评估模式的模型，供后续推理使用。
    return model

# model = init_model(device)
if __name__ == '__main__':
   model_path = r"weights/plate_rec_color.pth"
   image_path ="images/tmp2424.png"
   testPath = r"/mnt/Gpan/Mydata/pytorchPorject/CRNN/crnn_plate_recognition/images"
   fileList=[]
   allFilePath(testPath,fileList)
#    result = get_plate_result(image_path,device)
#    print(result)
   is_color = False
   model = init_model(device,model_path,is_color=is_color)
   right=0
   begin = time.time()
   
   for imge_path in fileList:
        img=cv2.imread(imge_path)
        if is_color:
            plate,_,plate_color,_=get_plate_result(img,device,model,is_color=is_color)
            print(plate)
        else:
            plate,_=get_plate_result(img,device,model,is_color=is_color)
            print(plate,imge_path)
        
  
        
